package com.zhang.二叉树;

import java.util.Comparator;

/**
 * @author 张吉鑫
 * @Description 平衡二叉树
 * @date 2021/3/13 21:33
 */
public class AVLTree<E> extends BBST<E>{
    public AVLTree() {
    }

    public AVLTree(Comparator comparator) {
        super(comparator);
    }


    /**
     * 添加后调整平衡因子
     * @param node 新添加的节点
     */
    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null){
            if (isBalanced(node)){
                //更新高度
                updateHeight(node);
            }else {
                //恢复平衡
                rebalance(node);
                break;
            }
        }
    }

    /**
     * 删除后调整平衡因子
     * @param node 替换被删除的节点
     */
    @Override
    protected void afterRemove(Node<E> node , Node<E> childNode) {
        while ((node = node.parent) != null){
            if (isBalanced(node)){
                updateHeight(node);
            }else {
                rebalance(node);
            }
        }
    }

    /**
     * 恢复平衡
     * @param grand 高度最低的那个不平衡节点
     */
    private void rebalance(Node<E> grand){
        Node<E> parent = ((AVLNode<E>)grand).tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();
        if (parent.isLeftChild()){//L
            if (node.isLeftChild()){//LL
                uniformRotation(grand, node.left, node, node.right, parent, parent.right, grand, grand.right);
                //rotateRight(grand);
            }else {//LR
                uniformRotation(grand, parent.left, parent, node.left, node, node.right, grand, grand.right);
//                rotateLeft(parent);
//                rotateRight(grand);
            }
        }else {//R
            if (node.isLeftChild()){//RL
                uniformRotation(grand, grand.left, grand, node.left, node, node.right, parent, parent.right);
//                rotateRight(parent);
//                rotateLeft(grand);
            }else {//RR
                uniformRotation(grand, grand.left, grand, parent.left, parent, node.left, node, node.right);
//                rotateLeft(grand);
            }
        }
    }

    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        super.afterRotate(grand, parent, child);
        //更新高度
        updateHeight(grand);
        updateHeight(parent);
    }

    /**
     * 统一选装操作
     * @param parent
     * @param a
     * @param b
     * @param c
     * @param d
     * @param e
     * @param f
     * @param g
     */
    @Override
    protected void uniformRotation(Node<E> parent, Node<E> a, Node<E> b, Node<E> c, Node<E> d, Node<E> e, Node<E> f, Node<E> g) {
        super.uniformRotation(parent, a, b, c, d, e, f, g);
        //更新高度
        updateHeight(b);
        updateHeight(f);
        updateHeight(d);
    }

    /**
     * 判断是否平衡
     * @param node
     * @return
     */
    private boolean isBalanced(Node<E> node){
        return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }

    /**
     * 更新高度
     * @param node
     */
    private void updateHeight(Node<E> node){
        ((AVLNode<E>)node).updateHeight();
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }


    /**
     * AVL节点元素
     * @param <E> 存储的元素
     */
    private static class AVLNode<E> extends Node<E>{
        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 平衡因子
         * @return 返回左右子节点的高度差 绝对值大于1 代表失去平衡
         */
        public int balanceFactor(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return leftHeight - rightHeight;
        }

        /**
         * 更新树的高度
         */
        public void updateHeight(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        /**
         * @return 获取左右子节点中最高的一边的子树。如果左右两边子树一样高，就返回祖子树顺方向的子树
         */
        public Node<E> tallerChild(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            if (leftHeight > rightHeight){
                return left;
            }
            if (leftHeight < rightHeight){
                return right;
            }
            /**
             * 如果左右两边子树一样高 就返回祖子树顺方向的子树
             */
            return isLeftChild() ? left : right;
        }

        @Override
        public String toString() {
            String parentString = "null";
            if (parent != null){
                parentString = parent.element.toString();
            }
            return element + "_p(" + parentString + ")_h(" + height + ")";
        }

    }
}
